#!/bin/bash

# NRC Rust编译器环境初始化脚本
# 检查并安装必要的编译环境

set -e

echo "🚀 NRC Rust编译器环境初始化"
echo "=================================="
echo ""

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 修复GPG密钥问题
fix_gpg_keys() {
    echo -e "${BLUE}检查并修复GPG密钥...${NC}"
    # 尝试修复Cursor的GPG密钥（使用现代方法）
    local key_id="42A1772E62E492D6"
    if [ -f "/etc/apt/sources.list.d/cursor.list" ] || [ -f "/etc/apt/sources.list.d/cursor.list.save" ]; then
        # 使用gpg --dearmor方式添加密钥（现代方法）
        if curl -fsSL https://downloads.cursor.com/aptrepo/gpg.key 2>/dev/null | sudo gpg --dearmor -o /usr/share/keyrings/cursor-archive-keyring.gpg 2>/dev/null; then
            echo -e "${GREEN}✅ GPG密钥修复成功${NC}"
            return 0
        fi
        # 备用方法：直接添加密钥
        if gpg --no-default-keyring --keyring /tmp/cursor-keyring.gpg --keyserver keyserver.ubuntu.com --recv-keys "$key_id" 2>/dev/null; then
            gpg --no-default-keyring --keyring /tmp/cursor-keyring.gpg --export "$key_id" | sudo gpg --dearmor -o /usr/share/keyrings/cursor-archive-keyring.gpg 2>/dev/null
            rm -f /tmp/cursor-keyring.gpg
            echo -e "${GREEN}✅ GPG密钥修复成功${NC}"
            return 0
        fi
    fi
    echo -e "${YELLOW}⚠️  GPG密钥修复失败，但不会影响主要软件包安装${NC}"
    return 1
}

# 安全的apt update（忽略GPG警告）
safe_apt_update() {
    echo -e "${BLUE}更新软件包列表...${NC}"
    # 先尝试修复GPG密钥
    fix_gpg_keys
    # 执行apt update，即使有GPG警告也继续（因为主要仓库可能已经更新成功）
    # 使用set +e临时禁用错误退出，然后恢复
    set +e
    sudo apt update 2>&1 | grep -v "NO_PUBKEY" | grep -v "没有数字签名" | grep -v "无法安全地用该源进行更新" || true
    local update_status=$?
    set -e
    # 即使有GPG警告，也认为更新成功（因为主要仓库通常已经更新）
    return 0
}

# 检查函数
check_command() {
    local cmd=$1
    local package=$2
    local install_cmd=$3
    
    if command -v "$cmd" &> /dev/null; then
        echo -e "${GREEN}✅ $cmd 已安装${NC}"
        return 0
    else
        echo -e "${RED}❌ $cmd 未安装${NC}"
        echo -e "${YELLOW}需要安装: $package${NC}"
        echo -e "${BLUE}安装命令: $install_cmd${NC}"
        return 1
    fi
}

# 安装函数
install_package() {
    local package=$1
    local install_cmd=$2
    
    echo ""
    echo -e "${YELLOW}是否安装 $package? (Y/n)${NC}"
    read -r response
    if [[ "$response" =~ ^[Yy]$ ]] || [[ -z "$response" ]]; then
        echo -e "${BLUE}正在安装 $package...${NC}"
        # 如果是apt update命令，使用安全的更新方式
        if [[ "$install_cmd" == *"apt update"* ]]; then
            # 先执行安全的apt update
            safe_apt_update
            # 提取安装命令（去掉"sudo apt update && "部分）
            local install_only="${install_cmd#sudo apt update && }"
            # 执行安装命令
            set +e
            eval "$install_only"
            local install_status=$?
            set -e
            if [ $install_status -eq 0 ]; then
                echo -e "${GREEN}✅ $package 安装成功${NC}"
            else
                echo -e "${RED}❌ $package 安装失败${NC}"
                exit 1
            fi
        else
            set +e
            eval "$install_cmd"
            local install_status=$?
            set -e
            if [ $install_status -eq 0 ]; then
                echo -e "${GREEN}✅ $package 安装成功${NC}"
            else
                echo -e "${RED}❌ $package 安装失败${NC}"
                exit 1
            fi
        fi
    else
        echo -e "${YELLOW}跳过 $package 安装${NC}"
        return 1
    fi
}

# 检查Rust
echo "📋 检查Rust环境..."
if check_command "rustc" "Rust" "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh"; then
    echo -e "${GREEN}   Rust版本: $(rustc --version)${NC}"
else
    install_package "Rust" "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y"
    echo -e "${BLUE}请重新加载shell环境: source ~/.cargo/env${NC}"
fi

echo ""

# 检查LLVM
echo "📋 检查LLVM环境..."
llvm_installed=false

if check_command "llvm-config-15" "LLVM 15" "sudo apt update && sudo apt install -y llvm-15-dev llvm-15-tools"; then
    echo -e "${GREEN}   LLVM版本: $(llvm-config-15 --version)${NC}"
    llvm_installed=true
elif check_command "llvm-config" "LLVM" "sudo apt update && sudo apt install -y llvm-dev llvm-tools"; then
    echo -e "${GREEN}   LLVM版本: $(llvm-config --version)${NC}"
    llvm_installed=true
fi

if [ "$llvm_installed" = false ]; then
    install_package "LLVM 15" "sudo apt update && sudo apt install -y llvm-15-dev llvm-15-tools"
fi

# 配置 LLVM 环境变量
echo "📋 配置 LLVM 环境变量..."
if command -v llvm-config-15 &> /dev/null; then
    llvm_prefix=$(llvm-config-15 --prefix)
    cargo_config_dir=".cargo"
    cargo_config_file="$cargo_config_dir/config.toml"
    
    # 创建 .cargo 目录
    mkdir -p "$cargo_config_dir"
    
    # 检查是否已存在配置
    if [ -f "$cargo_config_file" ] && grep -q "LLVM_SYS_150_PREFIX" "$cargo_config_file"; then
        echo -e "${GREEN}✅ LLVM 环境变量已配置${NC}"
    else
        # 添加或更新配置
        if [ -f "$cargo_config_file" ]; then
            # 如果文件存在但没有 LLVM 配置，添加配置
            if ! grep -q "\[env\]" "$cargo_config_file"; then
                echo "" >> "$cargo_config_file"
                echo "[env]" >> "$cargo_config_file"
            fi
            if ! grep -q "LLVM_SYS_150_PREFIX" "$cargo_config_file"; then
                echo 'LLVM_SYS_150_PREFIX = "'"$llvm_prefix"'"' >> "$cargo_config_file"
            fi
        else
            # 创建新配置文件
            cat > "$cargo_config_file" << EOF
[env]
# LLVM 15 路径配置
# 用于 llvm-sys crate 查找 LLVM 库
LLVM_SYS_150_PREFIX = "$llvm_prefix"
EOF
        fi
        echo -e "${GREEN}✅ LLVM 环境变量已配置: $llvm_prefix${NC}"
    fi
elif command -v llvm-config &> /dev/null; then
    llvm_version=$(llvm-config --version | cut -d. -f1)
    if [ "$llvm_version" = "15" ]; then
        llvm_prefix=$(llvm-config --prefix)
        cargo_config_dir=".cargo"
        cargo_config_file="$cargo_config_dir/config.toml"
        mkdir -p "$cargo_config_dir"
        cat > "$cargo_config_file" << EOF
[env]
# LLVM 15 路径配置
# 用于 llvm-sys crate 查找 LLVM 库
LLVM_SYS_150_PREFIX = "$llvm_prefix"
EOF
        echo -e "${GREEN}✅ LLVM 环境变量已配置: $llvm_prefix${NC}"
    fi
fi

echo ""

# 检查llc
echo "📋 检查LLVM编译器..."
if check_command "llc-15" "LLVM编译器" "sudo apt install -y llvm-15-tools"; then
    echo -e "${GREEN}   llc-15版本: $(llc-15 --version | head -n1)${NC}"
else
    install_package "LLVM编译器" "sudo apt install -y llvm-15-tools"
fi

echo ""

# 检查链接器
echo "📋 检查链接器..."
linker_installed=false

if check_command "gcc" "GCC编译器" "sudo apt update && sudo apt install -y gcc"; then
    echo -e "${GREEN}   GCC版本: $(gcc --version | head -n1)${NC}"
    linker_installed=true
fi

if [ "$linker_installed" = false ]; then
    install_package "GCC编译器" "sudo apt update && sudo apt install -y gcc"
fi

echo ""

# 检查构建工具
echo "📋 检查构建工具..."
if check_command "cargo" "Cargo" "已包含在Rust安装中"; then
    echo -e "${GREEN}   Cargo版本: $(cargo --version)${NC}"
else
    echo -e "${RED}❌ Cargo未找到，请重新安装Rust${NC}"
    exit 1
fi

echo ""

# 测试编译
echo "🧪 测试编译环境..."
if cargo build --release > /dev/null 2>&1; then
    echo -e "${GREEN}✅ 编译测试成功${NC}"
else
    echo -e "${RED}❌ 编译测试失败${NC}"
    echo -e "${YELLOW}请检查错误信息并重新运行此脚本${NC}"
    exit 1
fi

echo ""
echo "🎉 环境初始化完成！"
echo "=================================="
echo ""
echo "📚 使用方法:"
echo "  1. 构建项目: cargo build --release"
echo "  2. 运行测试: ./target/release/nrc test examples"
echo ""
echo "💡 提示: 现在可以直接使用编译器了！"
